# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: inkyu

import datetime
import logging

import tabulate

from coin.exchange.shared.kr_rest.currency import SharedCurrency
from coin.strategy.hk_hard_arb.products import NativeProducts
from coin.strategy.hk_hard_arb.util import PassAggExchangeMap


class FungibleArbProduct(object):
  def __init__(self, product, native_products):
    self.product = product
    self.native_products = native_products
    for native_product in native_products:
      assert product.symbol == native_product.symbol


class PassAggArbProduct(FungibleArbProduct):
  def __init__(self, *, shared, native_passive, native_aggressive):
    super().__init__(product=shared,
                     native_products=PassAggExchangeMap(passive=native_passive,
                                                        aggressive=native_aggressive))

  @property
  def native_passive(self):
    return self.native_products.passive

  @property
  def native_aggressive(self):
    return self.native_products.aggressive


class FungibleArbProducts(object):
  def __init__(self, quote_currency_str, products_str, ProductTypeList):
    self.shared = NativeProducts.FromSharedStrProduct(quote_currency_str, products_str)
    self.natives = (NativeProducts.FromSharedProduct(self.shared.quote_currency,
                                                     self.shared.products,
                                                     ProductType)
                    for ProductType in ProductTypeList)
    self._arb_product_list = [
        FungibleArbProduct(self.shared.products[idx],
                           tuple(native.products[idx]
                                 for native in self.natives))
        for idx in range(len(products_str))
    ]
    self._symbol_map = {symbol: idx for idx, symbol in enumerate(products_str)}

  @property
  def n_exchange(self):
    return len(self.natives)

  def by_idx(self, idx):
    return self._arb_product_list[idx]

  def by_symbol(self, symbol):
    return self.by_idx(self._symbol_map[symbol])

  @staticmethod
  def two(quote_currency_str, products_str, ProductType0, ProductType1):
    ret = FungibleArbProducts(quote_currency_str, products_str, [ProductType0, ProductType1])
    assert ret.n_exchange == 2
    return ret

  @staticmethod
  def multi(quote_currency_str, products_str, ProductTypeList):
    ret = FungibleArbProducts(quote_currency_str, products_str, ProductTypeList)
    return ret


class PassAggArbProducts(object):
  def __init__(self, quote_currency_str, products_str, ProductTypePassive, ProductTypeAggressive):
    self.shared = NativeProducts.FromSharedStrProduct(quote_currency_str, products_str)
    self.natives = PassAggExchangeMap(
        passive=NativeProducts.FromSharedProduct(quote_currency_shared=self.shared.quote_currency,
                                                 products_shared=self.shared.products,
                                                 ProductType=ProductTypePassive),
        aggressive=NativeProducts.FromSharedProduct(
            quote_currency_shared=self.shared.quote_currency,
            products_shared=self.shared.products,
            ProductType=ProductTypeAggressive),
    )
    self._symbol_map = {symbol: idx for idx, symbol in enumerate(products_str)}

  def by_idx(self, idx):
    return PassAggArbProduct(shared=self.shared.products[idx],
                             native_passive=self.natives.passive.products[idx],
                             native_aggressive=self.natives.aggressive.products[idx])

  def by_symbol(self, symbol):
    return self.by_idx(self._symbol_map[symbol])


class FungibleArbBalance(object):
  def __init__(self, *, base_balances, quote_balances, base_reserve):
    base_total = sum(base_balances)
    quote_total = sum(quote_balances)
    base_offset = base_total - base_reserve
    self.base_balances = base_balances
    self.quote_balances = quote_balances
    self.base_total = base_total
    self.quote_total = quote_total
    self.base_reserve = base_reserve
    self.base_offset = base_offset


class ArbInfoPrinter(object):
  def __init__(self, ogs, quote_currency, base_currencies):
    self._logger = logging.getLogger(__name__)

    self._ogs = ogs
    self._ogs_name = [og._og_info.CurrencyType._exchange for og in ogs]
    self._balance_currencies = [quote_currency] + base_currencies

  def print_balance(self):
    table = []
    for cur in self._balance_currencies:
      # use cur.currency, not cur, as ogs do not have same type
      bals = [og.get_balance(cur.currency) for og in self._ogs]
      table.append(['*', cur.currency, sum(bals)] + bals)

    self._logger.info(
        '----- Balance:\n* %s\n%s' %
        (str(datetime.datetime.now()),
         tabulate.tabulate(
             table, headers=['*', 'Currency', 'Total'] + self._ogs_name, floatfmt=".6f")))

  def get_balance(self, arb_product, reserve_balance):
    assert isinstance(arb_product, FungibleArbProduct)
    # use cur.currency, not cur, as ogs do not have same type
    base_bals = [og.get_balance(arb_product.product.base.currency) for og in self._ogs]
    quote_bals = [og.get_balance(arb_product.product.quote.currency) for og in self._ogs]

    return FungibleArbBalance(base_balances=base_bals,
                              quote_balances=quote_bals,
                              base_reserve=reserve_balance)


class ArbInfoPnlPrinter(object):
  def __init__(self, quote_currency, base_currencies):
    self._logger = logging.getLogger(__name__)

    self._quote_currency = quote_currency
    self._balance_currencies = [quote_currency] + base_currencies
    self._ogs = None
    self._reserves = {}
    self._prices = {}

  @property
  def ready(self):
    if self._ogs:
      return True
    else:
      return False

  def set_ogs(self, ogs):
    self._ogs = ogs
    self._ogs_name = [og._og_info.CurrencyType._exchange for og in ogs]

  def _process_product(self, currency, is_bnb=False):
    # use cur.currency, not cur, as ogs do not have same type
    if is_bnb:
      bals = [og.get_balances().get(currency.currency) for og in self._ogs]
      bals_sum = sum(bal for bal in bals if bal is not None)
    else:
      bals = [og.get_balance(currency.currency) for og in self._ogs]
      bals_sum = sum(bals)
    qty_quote = 0
    if currency == self._quote_currency:
      qty_quote = bals_sum
    reserve = self._reserves.get(currency, None)
    if is_bnb:
      imbalance = bals_sum - reserve if reserve else bals_sum
    else:
      imbalance = bals_sum - reserve if reserve else None

    amount = None
    estimated_base = None
    midp = self._prices.get(currency, None)
    if midp is not None and imbalance is not None:
      amount = midp * imbalance
      estimated_base = midp * bals_sum
    table_row = (['*', currency.currency] + bals
                 + [bals_sum, reserve, imbalance, amount, estimated_base])
    return (table_row, qty_quote, amount, estimated_base)

  def print_balance(self):
    table = []
    imbalance_total = 0
    qty_quote = 0
    estimated_base_total = 0

    for cur in self._balance_currencies:
      table_row, qty_quote_res, amount, estimated_base = self._process_product(cur)
      table.append(table_row)
      qty_quote += qty_quote_res
      if amount is not None:
        imbalance_total += amount
      if estimated_base is not None:
        estimated_base_total += estimated_base

    # BNB
    if any(og.exchange == 'Binance' for og in self._ogs):
      table_row, qty_quote_res, amount, estimated_base = self._process_product(
          SharedCurrency.FromStrCurrency('BNB'), is_bnb=True)
      table.append(table_row)
      qty_quote += qty_quote_res
      if amount is not None:
        imbalance_total += amount
      if estimated_base is not None:
        estimated_base_total += estimated_base

    pnl_str = 'Quote: %f, ImbalanceQuote: %f, AdjustedQuote: %f, EstimatedBase: %f' % (
        qty_quote, imbalance_total, qty_quote + imbalance_total, estimated_base_total)
    self._logger.info(
        '----- Balance:\n* %s\n%s\n%s' %
        (str(datetime.datetime.now()),
         tabulate.tabulate(
             table,
             headers=(['*', 'Currency'] + self._ogs_name
                      + ['Total', 'Reserve', 'Imbalance', 'ImbalanceAmount', 'EstimatedBase']),
             floatfmt=".6f"),
         pnl_str))

  def get_balance(self, arb_product, reserve_balance):
    assert isinstance(arb_product, FungibleArbProduct)
    # use cur.currency, not cur, as ogs do not have same type
    base_bals = [og.get_balance(arb_product.product.base.currency) for og in self._ogs]
    quote_bals = [og.get_balance(arb_product.product.quote.currency) for og in self._ogs]

    return FungibleArbBalance(base_balances=base_bals,
                              quote_balances=quote_bals,
                              base_reserve=reserve_balance)
